import rclpy
from rclpy.node import Node
from nav_msgs.msg import OccupancyGrid
from geometry_msgs.msg import PointStamped
from std_msgs.msg import Int8
import tf2_ros
from tf2_ros import TransformException
from tf2_ros.buffer import Buffer
from tf2_ros.transform_listener import TransformListener
from rclpy.qos import QoSProfile, QoSHistoryPolicy, QoSReliabilityPolicy

class ROSInterface:
    def __init__(self, node: Node, base_frame: str = 'base_link', map_frame: str = 'map'):
        """
        初始化ROS接口
        base_frame: 机器人基坐标系
        map_frame: 地图坐标系
        """
        self.node = node
        self.base_frame = base_frame
        self.map_frame = map_frame
        
        # TF2设置
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, node)
        
        # 创建发布者和订阅者
        qos = QoSProfile(
            reliability=QoSReliabilityPolicy.BEST_EFFORT,
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=1
        )
        
        self.map_pub = node.create_publisher(OccupancyGrid, 'map', 10)
        self.flag_sub = node.create_subscription(Int8, 'robot_flag', self._flag_callback, 10)
        
        # 回调函数
        self.flag_callback = None
        
    def set_flag_callback(self, callback):
        """设置标志回调函数"""
        self.flag_callback = callback
        
    def _flag_callback(self, msg: Int8):
        """标志消息回调"""
        if self.flag_callback:
            self.flag_callback(msg.data)
            
    def get_robot_position(self) -> tuple:
        """获取机器人在地图坐标系中的位置"""
        try:
            # 获取机器人在地图坐标系中的位置
            transform = self.tf_buffer.lookup_transform(
                self.map_frame,
                self.base_frame,
                rclpy.time.Time()
            )
            
            x = transform.transform.translation.x
            y = transform.transform.translation.y
            return x, y
        except TransformException as ex:
            self.node.get_logger().error(f'无法获取变换: {ex}')
            return None, None
            
    def publish_map(self, map_data: np.ndarray, origin: tuple, width: int, height: int, resolution: float):
        """
        发布地图
        
        Args:
            map_data: 地图数据数组
            origin: 地图原点 (x, y)
            width: 地图宽度 (像素)
            height: 地图高度 (像素)
            resolution: 地图分辨率 (米/像素)
        """
        # 创建OccupancyGrid消息
        map_msg = OccupancyGrid()
        map_msg.header.stamp = self.node.get_clock().now().to_msg()
        map_msg.header.frame_id = self.map_frame
        
        # 设置地图信息
        map_msg.info.resolution = resolution
        map_msg.info.width = width
        map_msg.info.height = height
        map_msg.info.origin.position.x = origin[0]
        map_msg.info.origin.position.y = origin[1]
        map_msg.info.origin.position.z = 0.0
        map_msg.info.origin.orientation.x = 0.0
        map_msg.info.origin.orientation.y = 0.0
        map_msg.info.origin.orientation.z = 0.0
        map_msg.info.origin.orientation.w = 1.0
        
        # 将地图数据展平为一维数组
        map_msg.data = map_data.flatten().tolist()
        
        # 发布地图
        self.map_pub.publish(map_msg)